"""
2D Poisson High-Frequency Experiment for RINN paper reproduction.
Solves -�u = f with high-frequency manufactured solution.
"""

import torch
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
from typing import Dict, Tuple, Optional, List
import time

# Add parent directory to path for imports
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from pielm.solver import create_pielm_solver, laplacian_operator
from rinn.rinn_es import create_rinn_es_solver
from common.utils import generate_collocation_points, generate_boundary_points, set_random_seed
from common.metrics import evaluate_model_comprehensive, print_metrics
from common.plotting import plot_comparison, plot_training_history, plot_convergence_study


class PoissonHighFreq2D:
    """
    2D Poisson equation with high-frequency manufactured solution.
    
    Problem: -�u = f in � = [0,1]�
             u = g on �
             
    Manufactured solution: u(x,y) = sin(k*�*x) * sin(k*�*y)
    where k is the frequency parameter.
    """
    
    def __init__(self, frequency: int = 4, domain: Optional[Dict] = None):
        """
        Initialize high-frequency Poisson problem.
        
        Args:
            frequency: Frequency parameter k
            domain: Domain specification (default: unit square)
        """
        self.frequency = frequency
        self.k = frequency
        
        if domain is None:
            self.domain = {'x_min': 0.0, 'x_max': 1.0, 'y_min': 0.0, 'y_max': 1.0}
        else:
            self.domain = domain
    
    def analytical_solution(self, points: torch.Tensor) -> torch.Tensor:
        """
        Compute analytical solution u(x,y) = sin(k�x) * sin(k�y).
        
        Args:
            points: Evaluation points (N, 2)
            
        Returns:
            torch.Tensor: Solution values (N, 1)
        """
        x = points[:, 0:1]
        y = points[:, 1:2]
        
        u = torch.sin(self.k * np.pi * x) * torch.sin(self.k * np.pi * y)
        return u
    
    def source_function(self, points: torch.Tensor) -> torch.Tensor:
        """
        Compute source function f = -�u = 2k���sin(k�x)sin(k�y).
        
        Args:
            points: Evaluation points (N, 2)
            
        Returns:
            torch.Tensor: Source values (N, 1)
        """
        x = points[:, 0:1]
        y = points[:, 1:2]
        
        f = 2 * (self.k * np.pi)**2 * torch.sin(self.k * np.pi * x) * torch.sin(self.k * np.pi * y)
        return f
    
    def boundary_values(self, points: torch.Tensor) -> torch.Tensor:
        """
        Compute boundary values (Dirichlet BC).
        
        Args:
            points: Boundary points (N, 2)
            
        Returns:
            torch.Tensor: Boundary values (N, 1)
        """
        return self.analytical_solution(points)
    
    def generate_training_data(self, 
                             n_interior: int = 1000,
                             n_boundary: int = 200,
                             method: str = 'random') -> Dict:
        """
        Generate training data for the problem.
        
        Args:
            n_interior: Number of interior collocation points
            n_boundary: Number of boundary points per edge
            method: Sampling method ('random', 'uniform', 'latin_hypercube')
            
        Returns:
            dict: Training data
        """
        # Interior points
        interior_points = generate_collocation_points(
            domain=self.domain,
            n_points=n_interior,
            method=method
        )
        
        # Boundary points
        boundary_points = generate_boundary_points(
            domain=self.domain,
            n_boundary=n_boundary // 4,  # Points per edge
            include_corners=True
        )
        
        # Compute values
        source_values = self.source_function(interior_points)
        boundary_values = self.boundary_values(boundary_points)
        
        return {
            'interior_points': interior_points,
            'boundary_points': boundary_points,
            'source_values': source_values,
            'boundary_values': boundary_values
        }
    
    def generate_test_data(self, n_test: int = 2500) -> Dict:
        """
        Generate test data for evaluation.
        
        Args:
            n_test: Number of test points
            
        Returns:
            dict: Test data
        """
        # Generate uniform test grid
        n_x = int(np.sqrt(n_test))
        n_y = n_test // n_x
        
        x = torch.linspace(self.domain['x_min'], self.domain['x_max'], n_x)
        y = torch.linspace(self.domain['y_min'], self.domain['y_max'], n_y)
        xx, yy = torch.meshgrid(x, y, indexing='ij')
        
        test_points = torch.stack([xx.flatten(), yy.flatten()], dim=1)
        true_solution = self.analytical_solution(test_points)
        
        return {
            'test_points': test_points,
            'true_solution': true_solution
        }


def run_single_experiment(problem: PoissonHighFreq2D,
                         method: str,
                         hidden_dim: int,
                         seed: int = 42,
                         training_params: Optional[Dict] = None) -> Dict:
    """
    Run single experiment with specified method.
    
    Args:
        problem: Poisson problem instance
        method: Method name ('PIELM', 'RINN-es')
        hidden_dim: Number of hidden neurons
        seed: Random seed
        training_params: Additional training parameters
        
    Returns:
        dict: Experiment results
    """
    set_random_seed(seed)
    
    # Generate training data
    train_data = problem.generate_training_data()
    test_data = problem.generate_test_data()
    
    print(f"\nRunning {method} with {hidden_dim} hidden neurons (seed={seed})")
    print(f"Training points: {len(train_data['interior_points'])} interior, "
          f"{len(train_data['boundary_points'])} boundary")
    
    start_time = time.time()
    
    if method == 'PIELM':
        # Standard PIELM
        solver = create_pielm_solver(
            input_dim=2,
            hidden_dim=hidden_dim,
            activation='tanh',
            seed=seed
        )
        
        training_info = solver.train(
            collocation_points=train_data['interior_points'],
            boundary_points=train_data['boundary_points'],
            pde_operator=laplacian_operator,
            pde_rhs=train_data['source_values'],
            boundary_values=train_data['boundary_values']
        )
        
        network = solver.network
        
    elif method == 'RINN-es':
        # RINN with early stopping
        default_params = {
            'max_epochs': 500,
            'print_every': 50
        }
        if training_params:
            default_params.update(training_params)
        
        solver = create_rinn_es_solver(
            input_dim=2,
            hidden_dim=hidden_dim,
            activation='tanh',
            epsilon=1e-2,
            learning_rate=1e-3,
            patience=50,
            seed=seed
        )
        
        training_info = solver.train(
            collocation_points=train_data['interior_points'],
            boundary_points=train_data['boundary_points'],
            pde_operator=laplacian_operator,
            pde_rhs=train_data['source_values'],
            boundary_values=train_data['boundary_values'],
            **default_params
        )
        
        network = solver.network
        
    else:
        raise ValueError(f"Unknown method: {method}")
    
    training_time = time.time() - start_time
    
    # Evaluate on test data
    metrics = evaluate_model_comprehensive(
        network=network,
        test_points=test_data['test_points'],
        true_solution=test_data['true_solution'],
        pde_operator=laplacian_operator,
        bc_points=train_data['boundary_points'],
        bc_values=train_data['boundary_values']
    )
    
    # Get predictions for visualization
    with torch.no_grad():
        predictions = network(test_data['test_points'])
    
    results = {
        'method': method,
        'hidden_dim': hidden_dim,
        'seed': seed,
        'frequency': problem.frequency,
        'training_time': training_time,
        'training_info': training_info,
        'metrics': metrics,
        'test_data': test_data,
        'predictions': predictions,
        'training_data': train_data
    }
    
    print(f"Training completed in {training_time:.2f}s")
    print_metrics(metrics, f"{method} Results")
    
    return results


def run_convergence_study(frequency: int = 4,
                         hidden_dims: List[int] = [50, 100, 200, 400],
                         n_trials: int = 3,
                         save_results: bool = True) -> Dict:
    """
    Run convergence study comparing PIELM and RINN-es.
    
    Args:
        frequency: Frequency parameter for manufactured solution
        hidden_dims: List of hidden dimensions to test
        n_trials: Number of trials per configuration
        save_results: Whether to save results to file
        
    Returns:
        dict: Convergence study results
    """
    print(f"Starting convergence study for 2D Poisson (frequency={frequency})")
    print(f"Hidden dimensions: {hidden_dims}")
    print(f"Trials per configuration: {n_trials}")
    
    problem = PoissonHighFreq2D(frequency=frequency)
    methods = ['PIELM', 'RINN-es']
    
    results = {method: [] for method in methods}
    
    for hidden_dim in hidden_dims:
        print(f"\n{'='*60}")
        print(f"Testing hidden_dim = {hidden_dim}")
        print('='*60)
        
        for method in methods:
            method_errors = []
            
            for trial in range(n_trials):
                seed = trial * 42 + hidden_dim  # Reproducible but different seeds
                
                try:
                    result = run_single_experiment(
                        problem=problem,
                        method=method,
                        hidden_dim=hidden_dim,
                        seed=seed
                    )
                    
                    error = result['metrics']['rel_l2_error']
                    method_errors.append(error)
                    
                except Exception as e:
                    print(f"Error in {method} trial {trial}: {str(e)}")
                    method_errors.append(float('inf'))
            
            # Store results
            valid_errors = [e for e in method_errors if e != float('inf')]
            if valid_errors:
                mean_error = np.mean(valid_errors)
                std_error = np.std(valid_errors)
                min_error = np.min(valid_errors)
                success_rate = len(valid_errors) / len(method_errors)
            else:
                mean_error = float('inf')
                std_error = 0.0
                min_error = float('inf')
                success_rate = 0.0
            
            results[method].append({
                'hidden_dim': hidden_dim,
                'mean_error': mean_error,
                'std_error': std_error,
                'min_error': min_error,
                'success_rate': success_rate,
                'all_errors': method_errors
            })
            
            print(f"{method:10s}: Mean L2 Error = {mean_error:.6e} � {std_error:.6e} "
                  f"(Success: {success_rate:.1%})")
    
    # Summary
    print(f"\n{'='*60}")
    print("CONVERGENCE STUDY SUMMARY")
    print('='*60)
    
    for method in methods:
        print(f"\n{method}:")
        for result in results[method]:
            hd = result['hidden_dim']
            me = result['mean_error']
            se = result['std_error']
            sr = result['success_rate']
            print(f"  Hidden={hd:3d}: {me:.6e} � {se:.6e} (Success: {sr:.1%})")
    
    # Plot convergence
    plot_convergence_results(results, frequency)
    
    if save_results:
        # Save to file
        import json
        results_file = f"poisson_2d_freq{frequency}_convergence.json"
        
        # Convert to JSON-serializable format
        json_results = {}
        for method, method_results in results.items():
            json_results[method] = []
            for r in method_results:
                json_results[method].append({
                    'hidden_dim': r['hidden_dim'],
                    'mean_error': float(r['mean_error']),
                    'std_error': float(r['std_error']),
                    'min_error': float(r['min_error']),
                    'success_rate': r['success_rate'],
                    'all_errors': [float(e) for e in r['all_errors']]
                })
        
        with open(results_file, 'w') as f:
            json.dump(json_results, f, indent=2)
        
        print(f"\nResults saved to {results_file}")
    
    return results


def plot_convergence_results(results: Dict, frequency: int):
    """Plot convergence study results."""
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    
    # Plot mean errors
    for method, method_results in results.items():
        hidden_dims = [r['hidden_dim'] for r in method_results]
        mean_errors = [r['mean_error'] for r in method_results]
        std_errors = [r['std_error'] for r in method_results]
        
        # Filter out infinite errors for plotting
        valid_indices = [i for i, e in enumerate(mean_errors) if e != float('inf')]
        if valid_indices:
            hd_valid = [hidden_dims[i] for i in valid_indices]
            me_valid = [mean_errors[i] for i in valid_indices]
            se_valid = [std_errors[i] for i in valid_indices]
            
            axes[0].errorbar(hd_valid, me_valid, yerr=se_valid, 
                           marker='o', label=method, linewidth=2, markersize=6)
    
    axes[0].set_xlabel('Number of Hidden Neurons')
    axes[0].set_ylabel('Relative L2 Error')
    axes[0].set_title(f'Convergence Study (Frequency={frequency})')
    axes[0].set_yscale('log')
    axes[0].grid(True, alpha=0.3)
    axes[0].legend()
    
    # Plot success rates
    for method, method_results in results.items():
        hidden_dims = [r['hidden_dim'] for r in method_results]
        success_rates = [r['success_rate'] for r in method_results]
        
        axes[1].plot(hidden_dims, success_rates, marker='s', 
                    label=method, linewidth=2, markersize=6)
    
    axes[1].set_xlabel('Number of Hidden Neurons')
    axes[1].set_ylabel('Success Rate')
    axes[1].set_title(f'Success Rate (Frequency={frequency})')
    axes[1].set_ylim(0, 1.1)
    axes[1].grid(True, alpha=0.3)
    axes[1].legend()
    
    plt.tight_layout()
    save_path = f'poisson_2d_freq{frequency}_convergence.png'
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    print(f"Convergence plot saved to {save_path}")
    plt.close()


def demonstrate_solution_quality(frequency: int = 4, hidden_dim: int = 200):
    """Demonstrate solution quality with visualization."""
    print(f"\nDemonstrating solution quality (frequency={frequency}, hidden_dim={hidden_dim})")
    
    problem = PoissonHighFreq2D(frequency=frequency)
    
    # Run both methods
    methods = ['PIELM', 'RINN-es']
    results = {}
    
    for method in methods:
        print(f"\nRunning {method}...")
        result = run_single_experiment(
            problem=problem,
            method=method,
            hidden_dim=hidden_dim,
            seed=42
        )
        results[method] = result
    
    # Create comparison plots
    for method, result in results.items():
        plot_comparison(
            points=result['test_data']['test_points'],
            true_values=result['test_data']['true_solution'],
            pred_values=result['predictions'],
            title=f"{method} - 2D Poisson (f={frequency})",
            save_path=f'poisson_2d_{method.lower()}_f{frequency}.png'
        )
        
        # Plot training history for RINN-es
        if method == 'RINN-es' and 'training_history' in result['training_info']:
            plot_training_history(
                history=result['training_info']['training_history'],
                save_path=f'poisson_2d_rinn_es_training_f{frequency}.png'
            )


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='2D Poisson High-Frequency Experiment')
    parser.add_argument('--frequency', type=int, default=4, help='Frequency parameter')
    parser.add_argument('--hidden_dim', type=int, default=200, help='Hidden dimension for demo')
    parser.add_argument('--convergence', action='store_true', help='Run convergence study')
    parser.add_argument('--demo', action='store_true', help='Run solution quality demo')
    parser.add_argument('--n_trials', type=int, default=3, help='Number of trials for convergence study')
    
    args = parser.parse_args()
    
    if args.convergence:
        print("Running convergence study...")
        hidden_dims = [50, 100, 200, 400]
        run_convergence_study(
            frequency=args.frequency,
            hidden_dims=hidden_dims,
            n_trials=args.n_trials
        )
    
    if args.demo:
        print("Running solution quality demonstration...")
        demonstrate_solution_quality(
            frequency=args.frequency,
            hidden_dim=args.hidden_dim
        )
    
    if not args.convergence and not args.demo:
        print("No experiment specified. Use --convergence or --demo")
        print("Example usage:")
        print("  python poisson_2d_high_freq.py --demo")
        print("  python poisson_2d_high_freq.py --convergence --n_trials 5")